import numpy as np
import cv2
from .utils import *




class BBoxToPoint:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "optional": {
                "bbox": ("BBOX",),
                "box": ("BOX",),
            },
        }

    RETURN_TYPES = ("Point",)
    RETURN_NAMES = ("Points",)
    FUNCTION = "bbox_to_coordinates"
    CATEGORY = "🌟FJDH/Bbox tool"

    def bbox_to_coordinates(self, bbox=None, box=None):
        points = []

        if bbox:
            x_min, y_min, x_max, y_max = bbox
            points.extend([
                (int(x_min), int(y_min)),
                (int(x_max), int(y_min)),
                (int(x_min), int(y_max)),
                (int(x_max), int(y_max))
            ])

        if box:
            x_min, y_min, x_max, y_max = box
            points.extend([
                (int(x_min), int(y_min)),
                (int(x_max), int(y_min)),
                (int(x_min), int(y_max)),
                (int(x_max), int(y_max))
            ])

        return (points,)

class PointToBBox:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "points": ("Point",),
            },
        }

    RETURN_TYPES = ("BBOX", "BOX",)
    RETURN_NAMES = ("BBox", "Box",)
    FUNCTION = "coordinates_to_bbox"
    CATEGORY = "🌟FJDH/Bbox tool"

    def coordinates_to_bbox(self, points):
        x_values = [point[0] for point in points]
        y_values = [point[1] for point in points]

        x_min = min(x_values)
        y_min = min(y_values)
        x_max = max(x_values)
        y_max = max(y_values)

        bbox = [x_min, y_min, x_max, y_max]
        box = [x_min, y_min, x_max, y_max]

        return (bbox, box,)


class BBoxAreaFilter:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "bboxes": ("BBOX",),
                "threshold": ("FLOAT", {
                    "default": 0.5,
                    "min": 0.0,
                    "max": 1.0,
                    "step": 0.01,
                    "display": "number"
                }),
            },
        }

    RETURN_TYPES = ("BBOX",)
    RETURN_NAMES = ("Bboxes",)
    FUNCTION = "filter_bboxes"
    CATEGORY = "🌟FJDH/Bbox tool"

    def filter_bboxes(self, bboxes, threshold):
        if len(bboxes) == 1:
            return (bboxes,)

        areas = [(bbox[2] - bbox[0]) * (bbox[3] - bbox[1]) for bbox in bboxes]

        mean_area = np.mean(areas)
        std_area = np.std(areas)

        lower_bound = mean_area - threshold * std_area
        upper_bound = mean_area + threshold * std_area

        filtered_bboxes = [bbox for bbox, area in zip(bboxes, areas) if lower_bound <= area <= upper_bound]

        if not filtered_bboxes:
            return (bboxes,)

        return (filtered_bboxes,)


class MaskToBBox:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "mask": ("MASK",),
            },
        }

    RETURN_TYPES = ("BBOX", "Point")
    RETURN_NAMES = ("BBoxes", "CenterPoints")
    FUNCTION = "convert_mask_to_bbox"
    CATEGORY = "🌟FJDH/Bbox tool"

    def convert_mask_to_bbox(self, mask):
        all_bboxes = []
        all_centers = []

        for batch_mask in mask:
            mask_np = batch_mask.squeeze().cpu().numpy()
            mask_binary = (mask_np > 0).astype(np.uint8)

            # 找到轮廓
            contours, _ = cv2.findContours(mask_binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

            for contour in contours:
                # 最小外接矩形
                x, y, w, h = cv2.boundingRect(contour)
                bbox = (x, y, x + w, y + h)
                all_bboxes.append(bbox)

                # 计算bbox中心点
                center_x = (bbox[0] + bbox[2]) // 2
                center_y = (bbox[1] + bbox[3]) // 2
                all_centers.append((center_x, center_y))

        return (all_bboxes, [all_centers])

class ItemSelector:
    @classmethod
    def INPUT_TYPES(s):
        return {
            "required": {
                "input_list": (any_type,),
                "index": ("INT", {"default": 0, "min": 0, "max": 1000000, "step": 1, "display": "number"}),
            },
        }

    RETURN_TYPES = (any_type, any_type, any_type)
    RETURN_NAMES = ("selected_item", "selected_item_without_brackets", "input_list_without_brackets")
    FUNCTION = "select_item_by_index"
    CATEGORY = "🌟FJDH/Tool"

    def select_item_by_index(self, input_list, index):
        # 如果输入是列表且索引在范围内，返回指定项
        if isinstance(input_list, list) and 0 <= index < len(input_list):
            selected_item = input_list[index]
            selected_item_without_brackets = input_list[index] if isinstance(input_list[index], (int, float, str)) else input_list[index][0]
            input_list_without_brackets = input_list[index] if isinstance(input_list[index], (int, float, str)) else input_list[index][0]
            return (selected_item, selected_item_without_brackets, input_list_without_brackets)
        # 如果输入是张量且索引在范围内，返回指定项
        elif isinstance(input_list, torch.Tensor) and 0 <= index < input_list.size(0):
            selected_item = input_list[index]
            selected_item_without_brackets = input_list[index] if input_list[index].dim() == 0 else input_list[index][0]
            input_list_without_brackets = input_list[index] if input_list[index].dim() == 0 else input_list[index][0]
            return (selected_item, selected_item_without_brackets, input_list_without_brackets)
        # 如果输入不是列表或张量，或者索引超出范围，返回输入本身
        return (input_list, input_list, input_list)


NODE_CLASS_MAPPINGS = {
    "BBoxToPoint": BBoxToPoint,
    "PointToBBox": PointToBBox,
    "BBoxAreaFilter": BBoxAreaFilter,
    "MaskToBBox": MaskToBBox,
    "ItemSelector": ItemSelector,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "BBoxToPoint": "BBox to Point",
    "PointToBBox": "Point to BBox",
    "BBoxAreaFilter": "BBox Area Filter",
    "MaskToBBox": "Mask to BBox",
    "ItemSelector": "Item Selector",
}